#ifndef _COMMUNICATION_H
#define _COMMUNICATION_H


#include <iostream>
#include <ros/ros.h>
#include <Eigen/Eigen>
#include "formation_avoidance/AstarPath.h"
#include "geometry_msgs/Point.h"
#include "formation_avoidance/OptimizeTraj.h"
#include "bspline_cius/non_uniform_bspline_cius.h"
#include "formation_avoidance/OptimizeTrajExchange.h"
#include<iosfwd>
#include<algorithm>
#include<stdio.h>
 #include<numeric>

namespace formation_avoidance {
class Communication
{
public:	
	Communication(int set_agent_num):loop_rate(400),agent_num(set_agent_num)
	{
		astar_path.resize(set_agent_num);
		agent_round.resize(set_agent_num);
		getAstar.resize(set_agent_num);
		send_current_state.resize(set_agent_num);
		local_path.resize(set_agent_num);
		agent_state.resize(set_agent_num);
		for(int id=0;id<set_agent_num;id++)
			agent_state[id].resize(set_agent_num);
		trajcost.resize(set_agent_num);
		gettraj.resize(set_agent_num);
		local_traj.resize(set_agent_num);

	}
	~Communication(){}
	void init(ros::NodeHandle &nh,int set_id,int set_agent_num)
	{
		communication_id=set_id;
		agent_num=set_agent_num;
		a_star_path_sub=nh.subscribe("/a_star_path", 5, &Communication::a_star_path_cb,this);
    	a_star_path_pub = nh.advertise<formation_avoidance::AstarPath>("/a_star_path", 1);

		optimize_traj_exchange_sub=nh.subscribe("/optimize_traj_exchange", 5, &Communication::optimize_traj_cb,this);

    	optimize_traj_exchange_pub = nh.advertise<formation_avoidance::OptimizeTrajExchange>("/optimize_traj_exchange", 1);


		round=0;
		for(int i=0;i<agent_num;i++)
			getAstar[i]=false;
		getAstar[communication_id]=true;
	}

	void a_star_path_cb(formation_avoidance::AstarPath msg)
	{
		if(msg.agentID==communication_id)return;
		if(getAstar[msg.agentID]) return;
		if(execflag>0.5)
		{
			if(msg.round!=round)return;
		}
		bestflag=msg.bestPath;
		if(bestflag) bestID=msg.agentID;
		astar_path[msg.agentID].clear();
		for(int i=0;i<(int)msg.path_point.size();i++)
		{
			Eigen::Vector3d point(msg.path_point[i].x,msg.path_point[i].y,msg.path_point[i].z);
			astar_path[msg.agentID].push_back(point);
		}
		getAstar[msg.agentID]=true;


		//store position and velocity from other agent

		for(int id=0;id<agent_num;id++)
		{
			agent_state[id][msg.agentID](0)=msg.current_agent_p[id].x;
			agent_state[id][msg.agentID](1)=msg.current_agent_p[id].y;
			agent_state[id][msg.agentID](2)=msg.current_agent_p[id].z;
			agent_state[id][msg.agentID](3)=msg.current_agent_v[id].x;
			agent_state[id][msg.agentID](4)=msg.current_agent_v[id].y;		
			agent_state[id][msg.agentID](5)=msg.current_agent_v[id].z;
		}
	}


	void send_Astar_path(std::vector<Eigen::Vector3d> path, bool bestPath)
	{
		formation_avoidance::AstarPath msg;
		msg.agentID=communication_id;
		msg.header.stamp=ros::Time::now();
		msg.round=round;
		msg.bestPath=bestPath;
		for(int i=0;i<(int)path.size();i++)
		{
			geometry_msgs::Point path_point;
			path_point.x=path[i](0);
			path_point.y=path[i](1);
			path_point.z=path[i](2);
			msg.path_point.push_back(path_point);
		}


		for(int id=0;id<agent_num;id++)
		{
			geometry_msgs::Point p;
			p.x=send_current_state[id](0);
			p.y=send_current_state[id](1);
			p.z=send_current_state[id](2);
			msg.current_agent_p.push_back(p);
			geometry_msgs::Point v;
			v.x=send_current_state[id](3);
			v.y=send_current_state[id](4);
			v.z=send_current_state[id](5);
			msg.current_agent_v.push_back(v);
		}


		a_star_path_pub.publish(msg);

	}


	void exchangeAstarPath(std::vector<Eigen::Vector3d>& path,std::vector<Eigen::Matrix<double,6,1>>& current_state)
	{
		round++;
		astar_path[communication_id]=path;
		for(int i=0;i<agent_num;i++)
			getAstar[i]=false;
		getAstar[communication_id]=true;
		bestflag=false;
		
		for(int id=0;id<agent_num;id++)
		{
			send_current_state[id]=current_state[id];
			agent_state[id][communication_id]=current_state[id];
		}
		double communication_start_time=ros::Time::now().toSec();
		while(1)
		{
			if((ros::Time::now().toSec()-communication_start_time)>5.0)
			{
				ROS_INFO("Id: %d can not exchange Astar path !!",communication_id);
				break;
			}
			if(bestflag)
			{
				path=astar_path[bestID];
				local_path=astar_path[bestID];
				for(int id=0;id<agent_num;id++)
				{
					current_state[id]=agent_state[id][bestID];
					send_current_state[id]=agent_state[id][bestID];
				}
				break;
			} 
			send_Astar_path(path,false);
			if(std::accumulate(getAstar.begin(),getAstar.end(),0)==agent_num )
			{
				Eigen::Vector3d centreEndPoint(0,0,0);
				for(int i=0;i<agent_num;i++)
				{
					centreEndPoint+=astar_path[i][astar_path[i].size()-1];
				}
				centreEndPoint/=agent_num;
				double tmp=0;
				int outlier=0;
				for(int i=0;i<agent_num;i++)
				{
					double distance=(astar_path[i][astar_path[i].size()-1]-centreEndPoint).norm();
					if(distance>tmp)
					{
						outlier=i;
						tmp=distance;
					}
				}
				path.clear();
				int choose_id=outlier==0 ? 1 : 0;
				path=astar_path[choose_id];
				local_path=astar_path[choose_id];

				for(int id=0;id<agent_num;id++)
				{
					current_state[id]=agent_state[id][choose_id];
					send_current_state[id]=agent_state[id][choose_id];
				}

				send_Astar_path(path,true);
				send_Astar_path(path,true);

				break;
			}
			ros::spinOnce();
			loop_rate.sleep();
		}

		for(int i=0;i<agent_num;i++)
			getAstar[i]=false;
		getAstar[communication_id]=true;
		bestflag=false;
	}

	void optimize_traj_cb(formation_avoidance::OptimizeTrajExchange msg)
	{
		if(msg.agentID==communication_id)return;
		if(gettraj[msg.agentID])return;
		if(msg.round!=traj_round)
		{
			if(msg.round==traj_round+1)
			{
				queryTrajFlag=true;
			}
			else{
				return;
			}
		}
		if(msg.cost<local_cost||queryTrajFlag)
		{
			local_cost=msg.cost;
			for(int id=0;id<agent_num;id++)
			{
				Eigen::MatrixXd control_points(msg.exchange_traj_control_points[id].control_points.size(),3);
				
				for(int i=0;i<msg.exchange_traj_control_points[id].control_points.size();i++)
				{
					geometry_msgs::Point point=msg.exchange_traj_control_points[id].control_points[i];
					control_points.row(i)<<point.x,point.y,point.z;
				}
				
				double interval=0.25;
				local_traj[id].setUniformBspline(control_points,3,interval);
				local_traj[id].start_time=msg.exchange_traj_control_points[id].traj_start_time.toSec();


				Eigen::VectorXd knot(msg.exchange_traj_control_points[id].knots.size());
				for(int i=0;i<msg.exchange_traj_control_points[id].knots.size();i++)
				{
					knot(i)=msg.exchange_traj_control_points[id].knots[i];
				}

				local_traj[id].setKnot(knot);
			}

		}
		gettraj[msg.agentID]=true;

	}


	void sendtraj(std::vector<trajectory_formation::NonUniformBspline> traj,double cost)
	{
		formation_avoidance::OptimizeTrajExchange msg;
		msg.agentID=communication_id;
		msg.cost=cost;
		for(int id=0;id<agent_num;id++)
		{
			formation_avoidance::OptimizeTraj send_traj;
			send_traj.traj_start_time=ros::Time(traj[id].start_time);
			for(int i=0;i<traj[id].control_points_.rows();i++)
			{
				geometry_msgs::Point pos_pts;
				pos_pts.x=traj[id].control_points_(i,0);
				pos_pts.y=traj[id].control_points_(i,1);
				pos_pts.z=traj[id].control_points_(i,2);
				send_traj.control_points.push_back(pos_pts);
			}
			for(int i=0;i<traj[id].u_.size();i++)
			{
				send_traj.knots.push_back(traj[id].u_(i));
			}
			msg.exchange_traj_control_points.push_back(send_traj);
		}
		msg.header.stamp=ros::Time::now();
		msg.round=traj_round;
		optimize_traj_exchange_pub.publish(msg);
	}

	void exchangeOptimizeTraj(std::vector<trajectory_formation::NonUniformBspline>& traj,double& cost)
	{

		queryTrajFlag=false;
		traj_round++;
		local_cost=cost;
		local_traj=traj;
		trajcost[communication_id]=cost;
		for(int i=0;i<agent_num;i++)
			gettraj[i]=false;
		gettraj[communication_id]=true;
		while(1)
		{
			sendtraj(local_traj,local_cost);
			ros::spinOnce();
			loop_rate.sleep();
			traj=local_traj;
			cost=local_cost;

			if(std::accumulate(gettraj.begin(),gettraj.end(),0)==agent_num  )
			{
				break;
			}
			
			send_Astar_path(local_path,true);


		}
		for(int i=0;i<agent_num;i++)
			gettraj[i]=false;
		gettraj[communication_id]=true;

	}

	bool queryTraj(std::vector<trajectory_formation::NonUniformBspline>& traj,double& cost)
	{
		for(int i=0;i<agent_num;i++)
			gettraj[i]=false;
		gettraj[communication_id]=true;
		queryTrajFlag=false;
		double start=ros::Time::now().toSec();
		while(1)
		{		
			loop_rate.sleep();
			ros::spinOnce();
			if(queryTrajFlag)
			{
				traj=local_traj;
				cost=local_cost;
				ROS_INFO("ID: %d query traj success!!!!",communication_id);		
				traj_round++;	
				break;
			}
			if(ros::Time::now().toSec()-start>0.002)
				break;
		}

		bool tmp=queryTrajFlag;
		queryTrajFlag=false;
		return tmp;
	}




public:
	int communication_id;
	ros::Subscriber a_star_path_sub;
	ros::Publisher a_star_path_pub;
	std::vector<std::vector<Eigen::Vector3d>> astar_path;
	std::vector<int> agent_round;
	int round;
	ros::Rate loop_rate;
	std::vector<bool> getAstar;
	bool bestflag=false;
	int execflag;
	std::vector<Eigen::Matrix<double, 6,1>> send_current_state;

	std::vector<Eigen::Vector3d> local_path;

	std::vector<std::vector<Eigen::Matrix<double, 6,1>>> agent_state;
	int bestID;

	std::vector<double> trajcost;

	ros::Publisher optimize_traj_exchange_pub;
	ros::Subscriber optimize_traj_exchange_sub;
	std::vector<bool> gettraj;

	double local_cost;
	std::vector<trajectory_formation::NonUniformBspline> local_traj ;
	int traj_round=0;

	bool queryTrajFlag=false;
	int agent_num;

};
}

#endif